package com.slipper.core.similarity;

import java.util.*;

public class CosineSimilarity {
    /**
     * 计算2个词频集合的相似度
     * @param vector1 词频集合1
     * @param vector2 词频集合2
     * @return 相似度
     */
    public static Double similarity(Map<String, Long> vector1, Map<String, Long> vector2) {
        if (vector1 == null || vector2 == null) {
            throw new IllegalArgumentException("词频不能为空！");
        }

        final Set<String> intersection = getIntersection(vector1, vector2);
        final double dotProduct = dot(vector1, vector2, intersection);
        double d1 = 0.0d;
        for (final Long value : vector1.values()) {
            d1 += Math.pow(value, 2);
        }
        double d2 = 0.0d;
        for (final Long value : vector2.values()) {
            d2 += Math.pow(value, 2);
        }

        if (d1 <= 0.0 || d2 <= 0.0) {
            return 0.0;
        } else {
            return dotProduct / (Math.sqrt(d1) * Math.sqrt(d2));
        }
    }

    /**
     * 获取2个词频集合的点积
     * @param vector1 词频集合1
     * @param vector2 词频集合2
     * @param intersection 词频集合交集
     * @return 点积
     */
    private static double dot(Map<String, Long> vector1, Map<String, Long> vector2, Set<String> intersection) {
        long dotProduct = 0;
        for (final String key : intersection) {
            dotProduct += vector1.get(key) * (long) vector2.get(key);
        }
        return dotProduct;
    }

    /**
     * 获取词频交集
     * @param vector1 词频集合1
     * @param vector2 词频集合2
     * @return 词频交集
     */
    private static Set<String> getIntersection(Map<String, Long> vector1, Map<String, Long> vector2) {
        final Set<String> intersection = new HashSet<>(vector1.keySet());
        intersection.retainAll(vector2.keySet());
        return intersection;
    }
}
